import os
import time

import albumentations as A
import cv2
import numpy as np
import torch
from PIL import Image
from albumentations.pytorch import ToTensorV2

from training import models
from training.datasets.transform import IsotropicResize
from training.tools.metrics import AverageMeter, ProgressMeter
from training.tools.train_utils import parse_args

mean = np.array((0.485, 0.456, 0.406), dtype=np.float32)
std = np.array((0.229, 0.224, 0.225), dtype=np.float32)


def generate_predict_mask(masks_pred, image_file, crop_image_name, args):
    mask_pred = masks_pred.squeeze().cpu().numpy()
    mask_pred_sum = np.sum(mask_pred)
    if mask_pred_sum > 0:
        print(mask_pred_sum)
    test_results_path = os.path.join(args.data_dir, 'test_crops_results')
    test_results_image_path = os.path.join(test_results_path, image_file[:-4])
    os.makedirs(test_results_image_path, exist_ok=True)
    # image = image * std * 255.0
    # image = image + mean * 255.0
    # image = Image.fromarray(image.astype(np.uint8)).resize(shape[::-1])
    # image_path = os.path.join(test_results_path, image_file['file'])
    # image.save(image_path)

    mask = Image.fromarray((mask_pred * 255).astype(np.uint8))
    mask_path = os.path.join(test_results_image_path, '{}.png'.format(crop_image_name[:-4]))
    mask.save(mask_path)


def predict_image_crops(args, model):
    transform = A.Compose([
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ToTensorV2(),
    ])
    model.eval()

    test_path = os.path.join(args.data_dir, 'test')
    test_crops_path = os.path.join(args.data_dir, 'test_crops')
    test_images = os.listdir(test_path)[:20]

    batch_time = AverageMeter('Time', ':6.3f')
    progress = ProgressMeter(len(test_images), [batch_time], prefix='Test: ')
    end = time.time()
    for i, image_file in enumerate(test_images):
        if not image_file.endswith('.jpg'):
            continue

        image = cv2.imread(os.path.join(test_path, image_file), cv2.IMREAD_COLOR)
        height, width = image.shape[0], image.shape[1]
        mask = np.zeros((height, width), dtype=np.int8)
        test_crops_image_path = os.path.join(test_crops_path, image_file[:-4])
        with torch.no_grad():
            for crop_image_name in os.listdir(test_crops_image_path):
                crop_image = cv2.imread(os.path.join(test_crops_image_path, crop_image_name))
                crop_image = cv2.cvtColor(crop_image, cv2.COLOR_BGR2RGB)
                transformed = transform(image=crop_image)
                crop_image = transformed["image"]
                crop_image = crop_image.unsqueeze(0).cuda()

                # compute output
                masks_pred = model(crop_image)
                masks_pred = torch.argmax(masks_pred, dim=1)

                # crop_image_name_split = crop_image_name[:-4].split('_')
                # row, col = int(crop_image_name_split[0]), int(crop_image_name_split[1])
                # row_end = min(height, (row + 1) * crop_size)
                # row_start = row_end - crop_size
                # col_end = min(width, (col + 1) * crop_size)
                # col_start = col_end - crop_size
                # mask[row_start:row_end, col_start:col_end] = masks_pred.squeeze().cpu().numpy()

                generate_predict_mask(masks_pred, image_file, crop_image_name, args)

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        if (i + 1) % args.print_freq == 0:
            progress.display(i + 1)


def predict_image(args, model):
    transform = A.Compose([
        IsotropicResize(max_side=512, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC),
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ToTensorV2(),
    ])
    model.eval()

    test_path = os.path.join(args.data_dir, 'test')
    test_images = os.listdir(test_path)

    batch_time = AverageMeter('Time', ':6.3f')
    progress = ProgressMeter(len(test_images), [batch_time], prefix='Test: ')
    end = time.time()
    with torch.no_grad():
        for i, image_file in enumerate(test_images):
            if not image_file.endswith('.jpg'):
                continue

            image = cv2.imread(os.path.join(test_path, image_file), cv2.IMREAD_COLOR)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            height, width = image.shape[0], image.shape[1]
            transformed = transform(image=image)
            image = transformed["image"]
            image = image.unsqueeze(0).cuda()

            # compute output
            masks_pred = model(image)
            masks_pred = torch.argmax(masks_pred, dim=1)

            mask_pred = masks_pred.squeeze().cpu().numpy()
            # mask_pred_sum = np.sum(mask_pred)
            # if mask_pred_sum > 0:
            #     print(mask_pred_sum)
            # test_results_path = os.path.join(args.data_dir, 'test_results')
            test_results_path = os.path.join(args.data_dir, 'images')

            mask = Image.fromarray((mask_pred * 255).astype(np.uint8)).resize((width, height))
            mask_path = os.path.join(test_results_path, '{}.png'.format(image_file[:-4]))
            mask.save(mask_path)

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()
            if (i + 1) % args.print_freq == 0:
                progress.display(i + 1)


def main():
    args = parse_args()
    print(args)
    os.makedirs(os.path.join(args.data_dir, 'test_crops_results'), exist_ok=True)
    os.makedirs(os.path.join(args.data_dir, 'test_results'), exist_ok=True)
    if args.resume:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

        print("Loading checkpoint '{}'".format(args.resume))
        model = models.__dict__[args.arch]()
        model.cuda()
        checkpoint = torch.load(args.resume, map_location="cpu")
        state_dict = checkpoint.get("state_dict", checkpoint)
        model.load_state_dict(state_dict, strict=True)

        print("Starting generation")
        # predict_image_crops(args, model)
        predict_image(args, model)


if __name__ == '__main__':
    main()
